import gc
import torch
import uvicorn
from fastapi import FastAPI, Request
import model_qwen2vl as model
import transformers
from pydantic import BaseModel

device = "cuda"  # the device to load the model onto
print('cuda.is_available()：' + str(torch.cuda.is_available()))
print('cuda.device_count()：' + str(torch.cuda.device_count()))
print('__version__：' + str(torch.__version__))
print('transformers：' + str(transformers.__version__))
print('cuda version' + str(torch.version.cuda))

app = FastAPI()


@app.get('/')
def index():
    return '易拍即识'


class RequestBody(BaseModel):
    image_url: str
    target_text: str

@app.post('/')
def handler(data: RequestBody):
    print(data)

    image_url = data.image_url
    if not image_url:
        return "参数：image_url不能为空"
    target_text = data.target_text
    if not target_text:
        return "参数：target_text不能为空"

    result = model.recognize(image_url, target_text)

    gc.collect()
    torch.cuda.empty_cache()

    print(result)
    return result

